
#include "scatter_elements_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <cstring>

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    ScatterElementsTilingData tiling;
    const gert::StorageShape* x1_shape = context->GetInputShape(0);
    int32_t data_sz = 1;
    for (int i = 0; i < x1_shape->GetStorageShape().GetDimNum(); i++)
        data_sz *= x1_shape->GetStorageShape().GetDim(i);
    //tiling.set_size(data_sz);
    uint32_t varlength = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
    uint32_t indiceslength = context->GetInputShape(1)->GetStorageShape().GetShapeSize();
    uint32_t updatelength = context->GetInputShape(2)->GetStorageShape().GetShapeSize();
    tiling.set_varlength(varlength);
    tiling.set_indiceslength(indiceslength);
    tiling.set_updatelength(updatelength);

    int32_t reduce = 0;
    const char *str = context->GetAttrs()->GetAttrPointer<char>(1);
    if (strcmp(str, "add") == 0) {
        reduce = 1;
    } else if (strcmp(str, "multiply") == 0) {
        reduce = 2;
    } else {
        reduce = 3;
    }   
    tiling.set_reduce(reduce);
    auto axisptr = context->GetAttrs()->GetInt(0);
    int32_t dimnum = context->GetInputShape(0)->GetStorageShape().GetDimNum();
    int32_t axis = *axisptr >= 0 ? *axisptr : dimnum + *axisptr;
    tiling.set_axis(axis);
    tiling.set_dimnum(dimnum);
    int32_t varshape[64];
    int32_t indexshape[64];
    int32_t updateshape[64];
    int32_t varshapeAccm[64];
    int32_t indexshapeAccm[64];
    int32_t updateshapeAccm[64];    
    int32_t varlastdimnum;
    int32_t indexlastdimnum;
    int32_t updatelastdimnum;
    if(reduce == 1 || reduce == 2) { 
        if(dimnum == 1) {
            context->SetTilingKey(100);
        } else {
            context->SetTilingKey(200);
            const gert::StorageShape* var_shape = context->GetInputShape(0);
            const gert::StorageShape* index_shape = context->GetInputShape(1);
            const gert::StorageShape* update_shape = context->GetInputShape(2);
            for(int32_t k = 0; k < dimnum; k++) {
                varshape[k] = var_shape->GetStorageShape().GetDim(k);
                indexshape[k] = index_shape->GetStorageShape().GetDim(k);
                updateshape[k] = update_shape->GetStorageShape().GetDim(k);
            }
            for(int32_t k = 0; k < dimnum; k++) {
                varshapeAccm[k] = 1;
                indexshapeAccm[k] = 1;
                updateshapeAccm[k] = 1;
                for(int32_t m = k + 1; m < dimnum; m++) {
                    varshapeAccm[k] *= varshape[m];
                    indexshapeAccm[k] *= indexshape[m];
                    updateshapeAccm[k] *= updateshape[m];
                }
                // std::cout << "varshapeAccm:" << varshapeAccm[k] << std::endl;
                // std::cout << "indexshapeAccm:" << indexshapeAccm[k] << std::endl;
                // std::cout << "updateshapeAccm:" << updateshapeAccm[k] << std::endl;
            }
        }
    }
    if(reduce == 3) {
        if(dimnum == 1) {
            context->SetTilingKey(300);
        } else {
            context->SetTilingKey(400);
            const gert::StorageShape* var_shape = context->GetInputShape(0);
            const gert::StorageShape* index_shape = context->GetInputShape(1);
            const gert::StorageShape* update_shape = context->GetInputShape(2);
            for(int32_t k = 0; k < dimnum; k++) {
                varshape[k] = var_shape->GetStorageShape().GetDim(k);
                indexshape[k] = index_shape->GetStorageShape().GetDim(k);
                updateshape[k] = update_shape->GetStorageShape().GetDim(k);
            }
            for(int32_t k = 0; k < dimnum; k++) {
                varshapeAccm[k] = 1;
                indexshapeAccm[k] = 1;
                updateshapeAccm[k] = 1;
                for(int32_t m = k + 1; m < dimnum; m++) {
                    varshapeAccm[k] *= varshape[m];
                    indexshapeAccm[k] *= indexshape[m];
                    updateshapeAccm[k] *= updateshape[m];
                }
                // std::cout << "varshapeAccm:" << varshapeAccm[k] << std::endl;
                // std::cout << "indexshapeAccm:" << indexshapeAccm[k] << std::endl;
                // std::cout << "updateshapeAccm:" << updateshapeAccm[k] << std::endl;
            }
            // if(axis == dimnum - 1) {
            //     context->SetTilingKey(410);
            //     varlastdimnum = var_shape->GetStorageShape().GetDim(axis);
            //     indexlastdimnum = index_shape->GetStorageShape().GetDim(axis);
            //     updatelastdimnum = update_shape->GetStorageShape().GetDim(axis);
            // }
        }
    }
    if(updatelength == 1) {
        context->SetTilingKey(500);
    }
    tiling.set_varshape(varshape);
    tiling.set_indexshape(indexshape);
    tiling.set_updateshape(updateshape);
    tiling.set_varshapeAccm(varshapeAccm);
    tiling.set_indexshapeAccm(indexshapeAccm);
    tiling.set_updateshapeAccm(updateshapeAccm);    
    tiling.set_varlastdimnum(varlastdimnum);
    tiling.set_indexlastdimnum(indexlastdimnum);
    tiling.set_updatelastdimnum(updatelastdimnum);
    

    std::cout << "tiling_key:" << context->GetTilingKey() << std::endl;
    std::cout << "varlength:" << varlength << std::endl;
    std::cout << "indiceslength:" << indiceslength << std::endl;
    std::cout << "updatelength:" << updatelength << std::endl;
    std::cout << "dimnum:" << dimnum << std::endl;
    std::cout << "reduce:" << reduce << std::endl;
    std::cout << "axis:" << axis << std::endl;
    std::cout << "varlastdimnum:" << varlastdimnum << std::endl;
    std::cout << "indexlastdimnum:" << indexlastdimnum << std::endl;
    std::cout << "updatelastdimnum:" << updatelastdimnum << std::endl;


    context->SetBlockDim(1);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    // const gert::Shape* x1_shape = context->GetInputShape(0);
    // gert::Shape* y_shape = context->GetOutputShape(0);
    // *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class ScatterElements : public OpDef {
public:
    explicit ScatterElements(const char* name) : OpDef(name)
    {
        this->Input("var")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT8, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("indices")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("updates")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT8, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("axis").AttrType(OPTIONAL).Int(0);
        this->Attr("reduce").AttrType(OPTIONAL).String("None");

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(ScatterElements);
}
